import streamlit as st
import akshare as ak
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from datetime import datetime, timedelta

# 缓存数据获取（优化版）
@st.cache_data
def get_stock_data(symbol, start, end):
    try:
        code = f"{symbol}"
        # 获取数据
        df = ak.stock_zh_a_hist(
            symbol=code, period="daily",
            start_date=start.strftime("%Y%m%d"),
            end_date=end.strftime("%Y%m%d"),
            adjust="qfq",
        )

        # 列名标准化
        df = df.rename(columns={
            '日期': 'date', '开盘': 'open', '收盘': 'close',
            '最高': 'high', '最低': 'low', '成交量': 'volume'
        })

        # 数据校验
        required_cols = ['date', 'open', 'close', 'high', 'low', 'volume']
        if not all(col in df.columns for col in required_cols):
            missing = set(required_cols) - set(df.columns)
            raise ValueError(f"缺失关键字段: {missing}")

        df['date'] = pd.to_datetime(df['date'])
        return df.set_index('date').sort_index()

    except Exception as e:
        st.error(f"数据获取失败：{str(e)}")
        st.info("常见原因：\n1. 股票代码需为6位数字\n2. 日期范围超出可获取范围\n3. 网络连接异常")
        return None

# 增强版KDJ计算（支持异常值过滤）
def calculate_kdj(df, n=9, m=3):
    df = df.copy()

    # 过滤异常值
    df = df[(df['high'] >= df['low']) &
            (df['close'].pct_change().abs() < 0.2)]  # 过滤涨跌幅>20%的异常数据

    # KDJ计算逻辑
    low_min = df['low'].rolling(n, min_periods=1).min()
    high_max = df['high'].rolling(n, min_periods=1).max()
    rsv = (df['close'] - low_min) / (high_max - low_min + 1e-6) * 100  # 避免除零
    df['K'] = rsv.ewm(alpha=1 / m).mean()
    df['D'] = df['K'].ewm(alpha=1 / m).mean()
    df['J'] = 3 * df['K'] - 2 * df['D']

    return df[['K', 'D', 'J']]

# 多周期共振检测（增强信号过滤）
def detect_resonance(main_df, params):
    # 计算日线KDJ
    daily_kdj = calculate_kdj(main_df, n=params['daily_n'], m=params['m'])

    # 计算周线KDJ（带数据对齐）
    weekly_df = main_df.resample('W-FRI').agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last'
    })
    weekly_kdj = calculate_kdj(weekly_df, n=params['weekly_n'], m=params['m'])
    weekly_kdj = weekly_kdj.resample('D').ffill().reindex(main_df.index, method='ffill')

    # 计算月线KDJ（带数据对齐）
    monthly_df = main_df.resample('M').agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last'
    })
    monthly_kdj = calculate_kdj(monthly_df, n=params['monthly_n'], m=params['m'])
    monthly_kdj = monthly_kdj.resample('D').ffill().reindex(main_df.index, method='ffill')

    # 合并数据
    merged = pd.concat([
        daily_kdj.add_prefix('D_'),
        weekly_kdj.add_prefix('W_'),
        monthly_kdj.add_prefix('M_')
    ], axis=1).ffill()

    # 信号生成逻辑
    merged['D_金叉'] = (merged['D_K'] > merged['D_D']) & (merged['D_K'].shift() <= merged['D_D'].shift())
    merged['W_金叉'] = (merged['W_K'] > merged['W_D']) & (merged['W_K'].shift() <= merged['W_D'].shift())
    merged['M_金叉'] = (merged['M_K'] > merged['M_D']) & (merged['M_K'].shift() <= merged['M_D'].shift())

    # 三线共振信号（叠加过滤条件）
    merged['共振信号'] = (
            merged['D_金叉'] &
            merged['W_金叉'] &
            merged['M_金叉'] &
            (merged['D_K'] < 50) &  # 日线处于低位
            (merged['W_J'] > merged['W_J'].shift(3)) &  # 周线J值趋势向上
            (merged['M_D'] < 60)  # 月线D值未超买
    )

    return merged[['D_K', 'D_D', 'D_J', 'D_金叉', 'W_金叉', 'M_金叉', '共振信号']]

# 主界面布局（增强可视化）
def app():
    st.title("KDJ多周期共振选股系统")

    # 侧边栏参数设置
    with st.sidebar:
        st.header("⚙️参数设置")
        code = st.text_input("股票代码（如：600519）", value="600519")
        start_date = st.date_input("开始日期", datetime.now() - timedelta(days=365))
        end_date = st.date_input("结束日期", datetime.now())

        st.subheader("KDJ参数")
        daily_n = st.slider("日线周期N", 3, 30, 5, help="建议小于周线周期")
        weekly_n = st.slider("周线周期N", 10, 60, 27, help="建议介于日线和月线之间")
        monthly_n = st.slider("月线周期N", 30, 120, 89, help="建议大于周线周期")
        m = st.slider("平滑周期M", 2, 5, 3, help="数值越小信号越敏感")

    # 获取数据
    df = get_stock_data(code, start_date, end_date)
    if df is None:
        return

    # 计算信号
    params = {'daily_n': daily_n, 'weekly_n': weekly_n, 'monthly_n': monthly_n, 'm': m}
    signals = detect_resonance(df, params)
    merged_df = df.join(signals).dropna()

    st.subheader(f"{code} KDJ共振分析图")

    # 创建双图布局
    fig = make_subplots(rows=2, cols=1, 
                       shared_xaxes=True,
                       vertical_spacing=0.05,
                       row_heights=[0.7, 0.3],
                       specs=[[{"secondary_y": True}], [{}]])

    # K线主图
    fig.add_trace(go.Candlestick(
        x=merged_df.index,
        open=merged_df['open'],
        high=merged_df['high'],
        low=merged_df['low'],
        close=merged_df['close'],
        name='K线',
        increasing_line_color='#FF4500',
        decreasing_line_color='#4682B4'
    ), row=1, col=1)

    # 信号标记配置
    signal_config = {
        'D_金叉': ('gold', 'circle', 1.02),
        'W_金叉': ('blue', 'square', 1.05),
        'M_金叉': ('purple', 'diamond', 1.08),
        '共振信号': ('lime', 'star', 1.12)
    }

    # 添加信号标记
    for signal, (color, symbol, offset) in signal_config.items():
        if merged_df[signal].any():
            dates = merged_df[merged_df[signal]].index
            fig.add_trace(go.Scatter(
                x=dates,
                y=merged_df.loc[dates, 'high'] * offset,
                mode='markers',
                marker=dict(
                    color=color,
                    size=12,
                    symbol=symbol,
                    line=dict(width=1, color='black')
                ),
                name=f'{signal}',
                hovertext=[f"收盘价：{merged_df.loc[d, 'close']:.2f}<br>日期：{d.strftime('%Y-%m-%d')}" 
                          for d in dates]
            ), row=1, col=1)

    # KDJ指标图
    for col, color, name in zip(['D_K', 'D_D', 'D_J'], 
                              ['#1E90FF', '#FF8C00', '#32CD32'],
                              ['K值', 'D值', 'J值']):
        fig.add_trace(go.Scatter(
            x=merged_df.index,
            y=merged_df[col],
            name=name,
            line=dict(color=color, width=1.5),
            hoverinfo='x+y+name'
        ), row=2, col=1)

    # 图表布局增强
    fig.update_layout(
        height=900,
        hovermode='x unified',
        legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0.3, xanchor="center"),
        margin=dict(l=20, r=20, t=40, b=20),
        xaxis=dict(rangeslider=dict(visible=False)),
        xaxis2=dict(title="日期", showgrid=True, gridwidth=0.5, gridcolor='#2d3a46'),
        yaxis=dict(title="价格", showgrid=True, gridwidth=0.5, gridcolor='#2d3a46'),
        yaxis2=dict(title="KDJ值", showgrid=True, gridwidth=0.5, gridcolor='#2d3a46')
    )

    # 隐藏不必要的坐标轴标签
    fig.update_xaxes(showticklabels=False, row=1, col=1)
    fig.update_xaxes(showticklabels=True, row=2, col=1)

    st.plotly_chart(fig, use_container_width=True)

if __name__ == "__main__":
    app()